
import re
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class XMLtoBIODataset:
    def __init__(self, data_file, output_file=None):
        """
        初始化数据集，用于解析XML格式的NER数据并生成BIO标签
        :param data_file: 输入数据文件路径
        :param output_file: 输出文件路径（可选）
        """
        self.data_file = data_file
        self.output_file = output_file
        self.data = self.load_data()
        if self.output_file:
            self.save_data()

    def load_data(self):
        """解析数据文件，生成句子、BIO标签和ID"""
        sentences, labels, ids = [], [], []
        try:
            with open(self.data_file, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    parts = line.split('\t')
                    if len(parts) < 2:
                        logger.warning(f"跳过无效行: {line[:50]}...")
                        continue
                    id_, text = parts[0], parts[1]
                    sentence, sentence_labels = self._parse_xml_text(text)
                    if sentence and len(sentence) == len(sentence_labels):
                        sentences.append(sentence)
                        labels.append(sentence_labels)
                        ids.append(id_)
                    else:
                        logger.warning(f"解析失败，句子或标签为空: {text[:50]}...")
            if not sentences:
                logger.error(f"从 {self.data_file} 未加载到有效句子！")
                raise ValueError(f"从 {self.data_file} 未加载到有效句子！")
            logger.info(f"从 {self.data_file} 加载了 {len(sentences)} 个样本")
            for i, (sentence, labels, id_) in enumerate(self.data[:3]):
                logger.info(f"样本 {i+1}: ID={id_}, 句子={' '.join(sentence[:10])}..., 标签={labels[:10]}...")
            return list(zip(sentences, labels, ids))
        except Exception as e:
            logger.error(f"加载 {self.data_file} 时出错: {str(e)}")
            raise

    def _parse_xml_text(self, text):
        """解析XML格式文本，生成BIO标签"""
        tokens, token_labels = [], []
        current_pos = 0
        pattern = r'<category="([^"]+)">([^<]+)</category>'
        matches = list(re.finditer(pattern, text))
        for match in matches:
            category = match.group(1)
            entity_text = match.group(2).strip()
            start, end = match.span()
            pre_text = text[current_pos:start].strip()
            if pre_text:
                pre_tokens = [t for t in re.split(r'(\s+|[.,!?;:\-\(\)])', pre_text) if t.strip() and t not in [' ', '\t']]
                tokens.extend(pre_tokens)
                token_labels.extend(['O'] * len(pre_tokens))
            entity_tokens = [t for t in re.split(r'(\s+|[.,!?;:\-\(\)])', entity_text) if t.strip() and t not in [' ', '\t']]
            if entity_tokens:
                tokens.append(entity_tokens[0])
                token_labels.append(f'B-{category}')
                for token in entity_tokens[1:]:
                    tokens.append(token)
                    token_labels.append(f'I-{category}' if token not in [',', '.', '!', '?', ';', ':', '-', '(', ')'] else 'O')
            current_pos = end
        post_text = text[current_pos:].strip()
        if post_text:
            post_tokens = [t for t in re.split(r'(\s+|[.,!?;:\-\(\)])', post_text) if t.strip() and t not in [' ', '\t']]
            tokens.extend(post_tokens)
            token_labels.extend(['O'] * len(post_tokens))
        return tokens, token_labels

    def save_data(self):
        """将解析后的句子、BIO标签和ID保存到文件"""
        try:
            with open(self.output_file, 'w', encoding='utf-8') as f:
                for sentence, labels, id_ in self.data:
                    sentence_str = ' '.join(sentence)
                    labels_str = ' '.join(labels)
                    f.write(f"{id_}\t{sentence_str}\t{labels_str}\n")
            logger.info(f"解析后的数据集已保存到 {self.output_file}")
        except Exception as e:
            logger.error(f"保存数据集到 {self.output_file} 时出错: {str(e)}")
            raise

if __name__ == "__main__":
    # 配置参数
    data_file = "/root/npl/data/NCBI_corpus_testing.txt"
    output_file = "/root/npl/data/parsed_bio_dataset_test.txt"

    # 解析数据集
    dataset = XMLtoBIODataset(data_file, output_file=output_file)

